{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/bin/sh: 1: pip: not found\n"
     ]
    }
   ],
   "source": [
    "! pip install scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'cubic_spline_planner'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-3-47e705095bd5>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 13\u001b[0;31m     \u001b[0;32mimport\u001b[0m \u001b[0mcubic_spline_planner\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     14\u001b[0m \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m     \u001b[0;32mraise\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'cubic_spline_planner'"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Path tracking simulation with LQR steering control and PID speed control.\n",
    "author Atsushi Sakai (@Atsushi_twi)\n",
    "\"\"\"\n",
    "import scipy.linalg as la\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "import numpy as np\n",
    "import sys\n",
    "sys.path.append(\"../../PathPlanning/CubicSpline/\")\n",
    "\n",
    "try:\n",
    "    import cubic_spline_planner\n",
    "except:\n",
    "    raise\n",
    "\n",
    "\n",
    "Kp = 1.0  # speed proportional gain\n",
    "\n",
    "# LQR parameter\n",
    "Q = np.eye(4)\n",
    "R = np.eye(1)\n",
    "\n",
    "# parameters\n",
    "dt = 0.1  # time tick[s]\n",
    "L = 0.5  # Wheel base of the vehicle [m]\n",
    "max_steer = np.deg2rad(45.0)  # maximum steering angle[rad]\n",
    "\n",
    "show_animation = True\n",
    "#  show_animation = False\n",
    "\n",
    "\n",
    "class State:\n",
    "\n",
    "    def __init__(self, x=0.0, y=0.0, yaw=0.0, v=0.0):\n",
    "        self.x = x\n",
    "        self.y = y\n",
    "        self.yaw = yaw\n",
    "        self.v = v\n",
    "\n",
    "\n",
    "def update(state, a, delta):\n",
    "\n",
    "    if delta >= max_steer:\n",
    "        delta = max_steer\n",
    "    if delta <= - max_steer:\n",
    "        delta = - max_steer\n",
    "\n",
    "    state.x = state.x + state.v * math.cos(state.yaw) * dt\n",
    "    state.y = state.y + state.v * math.sin(state.yaw) * dt\n",
    "    state.yaw = state.yaw + state.v / L * math.tan(delta) * dt\n",
    "    state.v = state.v + a * dt\n",
    "\n",
    "    return state\n",
    "\n",
    "\n",
    "def PIDControl(target, current):\n",
    "    a = Kp * (target - current)\n",
    "\n",
    "    return a\n",
    "\n",
    "\n",
    "def pi_2_pi(angle):\n",
    "    return (angle + math.pi) % (2 * math.pi) - math.pi\n",
    "\n",
    "\n",
    "def solve_DARE(A, B, Q, R):\n",
    "    \"\"\"\n",
    "    solve a discrete time_Algebraic Riccati equation (DARE)\n",
    "    \"\"\"\n",
    "    X = Q\n",
    "    maxiter = 150\n",
    "    eps = 0.01\n",
    "\n",
    "    for i in range(maxiter):\n",
    "        Xn = A.T @ X @ A - A.T @ X @ B @ \\\n",
    "            la.inv(R + B.T @ X @ B) @ B.T @ X @ A + Q\n",
    "        if (abs(Xn - X)).max() < eps:\n",
    "            break\n",
    "        X = Xn\n",
    "\n",
    "    return Xn\n",
    "\n",
    "\n",
    "def dlqr(A, B, Q, R):\n",
    "    \"\"\"Solve the discrete time lqr controller.\n",
    "    x[k+1] = A x[k] + B u[k]\n",
    "    cost = sum x[k].T*Q*x[k] + u[k].T*R*u[k]\n",
    "    # ref Bertsekas, p.151\n",
    "    \"\"\"\n",
    "\n",
    "    # first, try to solve the ricatti equation\n",
    "    X = solve_DARE(A, B, Q, R)\n",
    "\n",
    "    # compute the LQR gain\n",
    "    K = la.inv(B.T @ X @ B + R) @ (B.T @ X @ A)\n",
    "\n",
    "    eigVals, eigVecs = la.eig(A - B @ K)\n",
    "\n",
    "    return K, X, eigVals\n",
    "\n",
    "\n",
    "def lqr_steering_control(state, cx, cy, cyaw, ck, pe, pth_e):\n",
    "    ind, e = calc_nearest_index(state, cx, cy, cyaw)\n",
    "\n",
    "    k = ck[ind]\n",
    "    v = state.v\n",
    "    th_e = pi_2_pi(state.yaw - cyaw[ind])\n",
    "\n",
    "    A = np.zeros((4, 4))\n",
    "    A[0, 0] = 1.0\n",
    "    A[0, 1] = dt\n",
    "    A[1, 2] = v\n",
    "    A[2, 2] = 1.0\n",
    "    A[2, 3] = dt\n",
    "    # print(A)\n",
    "\n",
    "    B = np.zeros((4, 1))\n",
    "    B[3, 0] = v / L\n",
    "\n",
    "    K, _, _ = dlqr(A, B, Q, R)\n",
    "\n",
    "    x = np.zeros((4, 1))\n",
    "\n",
    "    x[0, 0] = e\n",
    "    x[1, 0] = (e - pe) / dt\n",
    "    x[2, 0] = th_e\n",
    "    x[3, 0] = (th_e - pth_e) / dt\n",
    "\n",
    "    ff = math.atan2(L * k, 1)\n",
    "    fb = pi_2_pi((-K @ x)[0, 0])\n",
    "\n",
    "    delta = ff + fb\n",
    "\n",
    "    return delta, ind, e, th_e\n",
    "\n",
    "\n",
    "def calc_nearest_index(state, cx, cy, cyaw):\n",
    "    dx = [state.x - icx for icx in cx]\n",
    "    dy = [state.y - icy for icy in cy]\n",
    "\n",
    "    d = [idx ** 2 + idy ** 2 for (idx, idy) in zip(dx, dy)]\n",
    "\n",
    "    mind = min(d)\n",
    "\n",
    "    ind = d.index(mind)\n",
    "\n",
    "    mind = math.sqrt(mind)\n",
    "\n",
    "    dxl = cx[ind] - state.x\n",
    "    dyl = cy[ind] - state.y\n",
    "\n",
    "    angle = pi_2_pi(cyaw[ind] - math.atan2(dyl, dxl))\n",
    "    if angle < 0:\n",
    "        mind *= -1\n",
    "\n",
    "    return ind, mind\n",
    "\n",
    "\n",
    "def closed_loop_prediction(cx, cy, cyaw, ck, speed_profile, goal):\n",
    "    T = 500.0  # max simulation time\n",
    "    goal_dis = 0.3\n",
    "    stop_speed = 0.05\n",
    "\n",
    "    state = State(x=-0.0, y=-0.0, yaw=0.0, v=0.0)\n",
    "\n",
    "    time = 0.0\n",
    "    x = [state.x]\n",
    "    y = [state.y]\n",
    "    yaw = [state.yaw]\n",
    "    v = [state.v]\n",
    "    t = [0.0]\n",
    "\n",
    "    e, e_th = 0.0, 0.0\n",
    "\n",
    "    while T >= time:\n",
    "        dl, target_ind, e, e_th = lqr_steering_control(\n",
    "            state, cx, cy, cyaw, ck, e, e_th)\n",
    "\n",
    "        ai = PIDControl(speed_profile[target_ind], state.v)\n",
    "        state = update(state, ai, dl)\n",
    "\n",
    "        if abs(state.v) <= stop_speed:\n",
    "            target_ind += 1\n",
    "\n",
    "        time = time + dt\n",
    "\n",
    "        # check goal\n",
    "        dx = state.x - goal[0]\n",
    "        dy = state.y - goal[1]\n",
    "        if math.sqrt(dx ** 2 + dy ** 2) <= goal_dis:\n",
    "            print(\"Goal\")\n",
    "            break\n",
    "\n",
    "        x.append(state.x)\n",
    "        y.append(state.y)\n",
    "        yaw.append(state.yaw)\n",
    "        v.append(state.v)\n",
    "        t.append(time)\n",
    "\n",
    "        if target_ind % 1 == 0 and show_animation:\n",
    "            plt.cla()\n",
    "            plt.plot(cx, cy, \"-r\", label=\"course\")\n",
    "            plt.plot(x, y, \"ob\", label=\"trajectory\")\n",
    "            plt.plot(cx[target_ind], cy[target_ind], \"xg\", label=\"target\")\n",
    "            plt.axis(\"equal\")\n",
    "            plt.grid(True)\n",
    "            plt.title(\"speed[km/h]:\" + str(round(state.v * 3.6, 2))\n",
    "                      + \",target index:\" + str(target_ind))\n",
    "            plt.pause(0.0001)\n",
    "\n",
    "    return t, x, y, yaw, v\n",
    "\n",
    "\n",
    "def calc_speed_profile(cx, cy, cyaw, target_speed):\n",
    "    speed_profile = [target_speed] * len(cx)\n",
    "\n",
    "    direction = 1.0\n",
    "\n",
    "    # Set stop point\n",
    "    for i in range(len(cx) - 1):\n",
    "        dyaw = abs(cyaw[i + 1] - cyaw[i])\n",
    "        switch = math.pi / 4.0 <= dyaw < math.pi / 2.0\n",
    "\n",
    "        if switch:\n",
    "            direction *= -1\n",
    "\n",
    "        if direction != 1.0:\n",
    "            speed_profile[i] = - target_speed\n",
    "        else:\n",
    "            speed_profile[i] = target_speed\n",
    "\n",
    "        if switch:\n",
    "            speed_profile[i] = 0.0\n",
    "\n",
    "    speed_profile[-1] = 0.0\n",
    "\n",
    "    return speed_profile\n",
    "\n",
    "\n",
    "def main():\n",
    "    print(\"LQR steering control tracking start!!\")\n",
    "    ax = [0.0, 6.0, 12.5, 10.0, 7.5, 3.0, -1.0]\n",
    "    ay = [0.0, -3.0, -5.0, 6.5, 3.0, 5.0, -2.0]\n",
    "    goal = [ax[-1], ay[-1]]\n",
    "\n",
    "    cx, cy, cyaw, ck, s = cubic_spline_planner.calc_spline_course(\n",
    "        ax, ay, ds=0.1)\n",
    "    target_speed = 10.0 / 3.6  # simulation parameter km/h -> m/s\n",
    "\n",
    "    sp = calc_speed_profile(cx, cy, cyaw, target_speed)\n",
    "\n",
    "    t, x, y, yaw, v = closed_loop_prediction(cx, cy, cyaw, ck, sp, goal)\n",
    "\n",
    "    if show_animation:  # pragma: no cover\n",
    "        plt.close()\n",
    "        plt.subplots(1)\n",
    "        plt.plot(ax, ay, \"xb\", label=\"input\")\n",
    "        plt.plot(cx, cy, \"-r\", label=\"spline\")\n",
    "        plt.plot(x, y, \"-g\", label=\"tracking\")\n",
    "        plt.grid(True)\n",
    "        plt.axis(\"equal\")\n",
    "        plt.xlabel(\"x[m]\")\n",
    "        plt.ylabel(\"y[m]\")\n",
    "        plt.legend()\n",
    "\n",
    "        plt.subplots(1)\n",
    "        plt.plot(s, [np.rad2deg(iyaw) for iyaw in cyaw], \"-r\", label=\"yaw\")\n",
    "        plt.grid(True)\n",
    "        plt.legend()\n",
    "        plt.xlabel(\"line length[m]\")\n",
    "        plt.ylabel(\"yaw angle[deg]\")\n",
    "\n",
    "        plt.subplots(1)\n",
    "        plt.plot(s, ck, \"-r\", label=\"curvature\")\n",
    "        plt.grid(True)\n",
    "        plt.legend()\n",
    "        plt.xlabel(\"line length[m]\")\n",
    "        plt.ylabel(\"curvature [1/m]\")\n",
    "\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
